import mindspore
from mindspore import nn
from get_dataset2model import train_iterator, src_pad_idx, trg_pad_idx, model


"""模型训练 & 模型评估
定义损失函数与优化器。

损失函数：定义如何计算模型输出(logits)与目标(targets)之间的误差，这里可以使用交叉熵损失（CrossEntropyLoss）
优化器：MindSpore将模型优化算法的实现称为优化器。优化器内部定义了模型的参数优化过程（即梯度如何更新至模型参数），所有优化逻辑都封装在优化器对象中。

"""
loss_fn = nn.CrossEntropyLoss(ignore_index=trg_pad_idx)  # ignore_index=trg_pad_idx 表示在计算损失时忽略填充标记。
optimizer = nn.Adam(model.trainable_params(), learning_rate=0.0001)

"""模型训练逻辑
1.Network+loss function直接构造正向函数
2.函数变换，获得梯度计算（反向传播）函数
3.构造训练过程函数
4.调用函数进行训练

定义前向网络计算逻辑。
    在训练过程中，表示句子结尾的 <eos> 占位符应是被模型预测出来，而不是作为模型的输入，所以在处理 Decoder 的输入时，我们需要移除目标序列最末的 <eos> 占位符。
        trg = [, x_1, x_2, ..., x_n, eos]
        trg[:-1] = [, x_1, x_2, ..., x_n]
        其中，x_i代表目标序列中第i个表示实际内容的词元。
    
    我们期望最终的输出包含表示句末的 <eos> ，不包含表示句首的 <bos>，所以在计算损失时，需要同样去除的目标序列的句首 <bos> 占位符，再进行比较。???
        output = [y_1, y_2, ..., y_n, <eos>]
        trg[1:] = [x_1, x_2, ..., x_n, <eos>]
        其中，y_i表示预测的第i个实际内容词元。 
"""


def forward(enc_inputs, dec_inputs):
    """前向网络
    enc_inputs: [batch_size, src_len]
    dec_inputs: [batch_size, trg_len]
    """
    logits, _, _, _ = model(enc_inputs, dec_inputs[:, :-1], src_pad_idx, trg_pad_idx)

    targets = dec_inputs[:, 1:].view(-1)  # 目标序列，去除句首 <bos> 占位符
    loss = loss_fn(logits, targets)

    return loss


# 由于使用Cell封装神经网络模型，模型参数为Cell的内部属性，此时我们不需要使用grad_position指定对函数输入求导，因此将其配置为None。
# 对模型参数求导时，我们使用weights参数，使用model.trainable_params()方法从Cell中取出可以求导的参数。

grad_fn = mindspore.value_and_grad(forward, None, optimizer.parameters)


# 定义训练一个step的逻辑。
def train_step(enc_inputs, dec_inputs):
    loss, grads = grad_fn(enc_inputs, dec_inputs)
    optimizer(grads)
    return loss


# #定义整体训练逻辑。
#
# 在训练中，模型会以最小化损失为目标更新模型权重，故模型状态需设置为训练model.set_train(True)。

from tqdm import tqdm


def train(iterator, epoch=0):
    model.set_train(True)
    num_batches = len(iterator)
    total_loss = 0
    total_steps = 0

    with tqdm(total=num_batches) as t:
        t.set_description(f'Epoch: {epoch}')
        for src, src_len, trg in iterator():
            loss = train_step(src, trg)
            total_loss += loss.asnumpy()
            total_steps += 1
            curr_loss = total_loss / total_steps
            t.set_postfix({'loss': f'{curr_loss:.2f}'})
            t.update(1)

    return total_loss / total_steps


# #定义模型评估逻辑。
#
# 在评估中，仅需正向计算loss，无需更新模型参数,故模型状态需设置为训练model.set_train(False)。


def evaluate(iterator):
    model.set_train(False)
    num_batches = len(iterator)
    total_loss = 0
    total_steps = 0

    with tqdm(total=num_batches) as t:
        for src, _, trg in iterator():
            loss = forward(src, trg)
            total_loss += loss.asnumpy()
            total_steps += 1
            curr_loss = total_loss / total_steps
            t.set_postfix({'loss': f'{curr_loss:.2f}'})
            t.update(1)

    return total_loss / total_steps


if __name__ == '__main__':

    # 模型训练
    # 数据集遍历迭代，一次完整的数据集遍历成为一个epoch。
    # 我们逐个epoch打印训练的损失值和评估精度，并通过save_checkpoint保存ckpt文件（transformer.ckpt）

    from mindspore import save_checkpoint

    num_epochs = 10
    best_valid_loss = float('inf')
    ckpt_file_name = './transformer_en2cn.ckpt'

    for i in range(num_epochs):
        train_loss = train(train_iterator, i)
        # valid_loss = evaluate(valid_iterator)

        if i % 2 == 0:
            save_checkpoint(model, ckpt_file_name)
